from flask import Flask, render_template, request, jsonify
import cv2
import os
from MMEdu import MMClassification as cls
from openai import OpenAI

app = Flask(__name__)

# 初始化分类模型
model = cls('LeNet')
checkpoint = 'model/epoch_1 (1).pth'
classes = ['黑色舌苔', '地图舌苔', '紫色舌苔', '舌红苔黄苔厚腻', '红舌厚腻', '白舌厚腻']

# 初始化OpenAI客户端
client = OpenAI(
    api_key="sk-MY0mcNXpZGxoqUPPilZdX11KdGO5bgJZoiky62Fa586Xj3CC",
    base_url="https://api.moonshot.cn/v1",
)

# 加载提示文本
with open("D:/高中/AI中医/prompt_relic_koze-Copy1.txt", "r", encoding="utf-8") as f:
    prompts = f.read()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/analyze', methods=['POST'])
def analyze():
    try:
        # 获取上传的图片文件
        image_file = request.files['image']
        temp_path = 'temp.jpg'
        image_file.save(temp_path)

        # 执行推理
        result = model.inference(image=temp_path, checkpoint=checkpoint)
        pred = model.print_result(result)[0]
        
        # 获取医学建议
        history = [{"role": "system", "content": prompts}]
        response = client.chat.completions.create(
            model="moonshot-v1-8k",
            messages=history + [{"role": "user", "content": f"请讲解{pred['预测结果']}的成因和建议"}],
            temperature=0.3
        )
        
        advice = response.choices[0].message.content

        return jsonify({
            'class': pred['预测结果'],
            'advice': advice
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True)